from dataclasses import dataclass
from typing import List, Type, Dict, Optional, Union
from emi.bar.data import SymbolMeta, CategoryMeta, CategorySymbolMeta, BarMeta
from emi.bar.storage.model import BaseDBModel
from emi.bar.db import MarketMeta
from emi.core.DBSource import DBSource
import peewee
from datetime import datetime

META_TYPE_SYMBOL = 0   ## 个股详情
META_TYPE_CATEGORY = 1  ## 分类
META_TYPE_CATEGORY_SYMBOL = 2  ## 分类成分股
META_TYPE_NAME =  3   ## 分类里的CategoryMeta.type值对应的name值。

@dataclass
class _CategoryTypeMeta(BarMeta):
    """
    分类类型里面的type对应的nam值。
    """
    name:str
    type:int = 0 ## 对应的type值。

class MarketMetaModel(BaseDBModel):
    """
    指数元数据对象模型。
    code,meta_type,parent_code组成唯一值。
    """
    code = peewee.CharField(index=True,max_length=64)
    meta_type:int = peewee.SmallIntegerField(index=True) ## 0 是IndexMeta, 1是CategoryMeta，2是
    category_code = peewee.CharField(index=True,max_length=64,null=False)
    type:int = peewee.SmallIntegerField(index=True)
    name:str = peewee.CharField(max_length=64)
    weight: int = peewee.FloatField(default=None, null=True)
    time1:datetime = peewee.TimestampField(null=True)
    time2: datetime = peewee.TimestampField(null=True)
    _int1: int = peewee.IntegerField(default=-1)
    _int2: int = peewee.IntegerField(default=-1)
    _int3: int = peewee.IntegerField(default=-1)
    _float1: int = peewee.FloatField(default=None, null=True)
    _float2: int = peewee.FloatField(default=None, null=True)
    _float3: int = peewee.FloatField(default=None, null=True)
    _str1: str = peewee.CharField(default=None, null=True, max_length=256)
    _str2: str = peewee.CharField(default=None, null=True, max_length=256)
    _str3: str = peewee.CharField(default=None, null=True, max_length=256)
    core_data: bytes = peewee.BlobField(null=True)

    @staticmethod
    def from_meta(meta: Union[SymbolMeta,CategoryMeta], model=None):
        """
        IndexMeta、CategoryMemberMeta,CategoryMeta 转换成IndexDataModel
        id复合键值： (("code","category_code","meta_type",), True),

        """
        if model is None:
            model = MarketMetaModel()
        if isinstance(meta,CategorySymbolMeta):
            model.code = meta.code
            model.meta_type = META_TYPE_CATEGORY_SYMBOL
            model.type = -1
            model.category_code = meta.category_code
            model.name = meta.name
            model.weight = meta.weight
        elif isinstance(meta,CategoryMeta):
            model.code = meta.code
            model.type = meta.type
            model.meta_type = META_TYPE_CATEGORY
            model.category_code = ""
            model.name = meta.name
        elif isinstance(meta,SymbolMeta):
            model.code = meta.code
            model.type = -1
            model.meta_type = META_TYPE_SYMBOL
            model.category_code = ""
            model.name = meta.name
        elif isinstance(meta,_CategoryTypeMeta):
            model.code = ""
            model.type = meta.type
            model.meta_type = META_TYPE_NAME
            model.category_code = ""
            model.name = meta.name
        else:
            raise NotImplementedError("")
        return model


    def to_meta(self)->Union[SymbolMeta,CategoryMeta,CategorySymbolMeta,_CategoryTypeMeta]:
        if self.meta_type == META_TYPE_CATEGORY:
            return CategoryMeta(code=self.code,name=self.name,type=self.type)

        elif self.meta_type == META_TYPE_CATEGORY_SYMBOL:
            return CategorySymbolMeta(code=self.code,
                                      name=self.name,
                                      category_code=self.category_code,
                                      weight=self.weight
                                      )
        elif self.meta_type == META_TYPE_SYMBOL:
            return SymbolMeta(code= self.code,
                              name = self.name)
        elif self.meta_type == META_TYPE_NAME:
            return _CategoryTypeMeta(type= self.type,
                              name = self.name)
        raise NotImplementedError("")


class MarketMetaStorage(MarketMeta):
    """
    指数元数据存储器。
    """

    def __init__(self,source:DBSource,table_name = 'market_meta',index_metas:List[SymbolMeta]=None):
        indexes_desc = (
            (("code", "category_code", "meta_type",), True),
        )
        self.db_source = source
        self.index_metas = [] if index_metas is None else index_metas
        self.source = source.init_peewee(table_name,MarketMetaModel,indexes=indexes_desc,ignore_inited=True)

    def get_index_metas(self)->List[SymbolMeta]:
        return self.index_metas

    def import_from(self,marekt):
        """
        清除当前的市场结构数据，并导入marketMeta的数据。
        :return:
        """
        marekt:MarketMetaStorage = marekt

        with marekt.source.use_peewee() as other_model:
            models = (
                other_model.select()
            )
            model_list = [ d for d in models]

        with self.source.use_peewee() as model:
            try:
                model.drop_table()
            except Exception:
                pass
            model.create_table()
            model.save_all(model_list)



    def set_category_type_name(self,type:int,name:str):
        meta = _CategoryTypeMeta(type=type,name=name)
        model_data = MarketMetaModel.from_meta(meta)
        with self.source.use_peewee() as db_model:
            db_model.save_all([model_data])

    def get_category_types(self)->List[int]:
        """
        返回categy.type对应的名称。
        :return:
        """
        with self.source.use_peewee() as db_model:
            s = (
                db_model.select(db_model.type)
                    .where(
                    (db_model.meta_type == META_TYPE_CATEGORY)
                ).group_by(db_model.type).distinct()
            )
            return [ model.type for model in s]

    def get_category_type_name_map(self)->Dict[int,str]:
        """
        返回categy.type对应的名称。
        :return:
        """
        types = self.get_category_types()
        if not types:
            return {}
        with self.source.use_peewee() as db_model:
            s = (
                db_model.select()
                    .where(
                    (db_model.meta_type == META_TYPE_NAME)
                )
            )
            ret = {}
            for _type in types:
                ret[_type] = str(_type)
            for model in s:
                meta:_CategoryTypeMeta = model.to_meta()
                if ret.get(meta.type) is not None:
                    ret[meta.type] = meta.name
        return ret

    def save_category(self, datas: List[Union[CategoryMeta, CategorySymbolMeta]]):
        """
        保存行业分类结构信息。
        :update_symbol_meta 是或同步个股结构信息。
        """
        ds = []
        new_category_symbol_metas = {}
        for data in datas:
            if not isinstance(data,CategorySymbolMeta) and not isinstance(data,CategoryMeta):
                raise RuntimeError("can't save_category():type error!")
            model = MarketMetaModel.from_meta(data)
            ds.append(model)
            if isinstance(data,CategorySymbolMeta):
                new_category_symbol_metas[data.code] = data

        with self.source.use_peewee() as db_model:
            db_model.save_all(ds)

        if new_category_symbol_metas:
            ## 有增加新的行业成员股
            self._on_add_category_symbol_meta(list(new_category_symbol_metas.values()))

    def update_symbol_meta(self,data:SymbolMeta):
        """
        更新个股结构信息。
        """
        with self.source.use_peewee() as db_model:
            s = (
                db_model.select()
                    .where(
                    (db_model.meta_type == META_TYPE_SYMBOL)
                    & (db_model.code == data.code)
                ).first()
            )
            if not s:
                raise RuntimeError(f"cant find SymboleMeta: code = {data.code}")
            else:
                ds = db_model.from_meta(data,s)
            db_model.save_all([ds])


    def _on_add_category_symbol_meta(self,metas:List[CategorySymbolMeta]):
        if not metas:
            return
        symbols = [ meta.code  for meta in metas]
        symbol_meta_map = self.get_symbol_metas(symbols,as_map=True)
        no_exist_list = []
        for meta in metas:
            if symbol_meta_map.get(meta.code) is None:
                no_exist_list.append(meta)
        ##自动添加新的SymbolMeta数据。
        if no_exist_list:
            new_symbole_models = []
            for meta in no_exist_list:
                symbol_meta = SymbolMeta(code=meta.code,name=meta.name)
                new_symbole_models.append(MarketMetaModel.from_meta(symbol_meta))
            with self.source.use_peewee() as db_model:
                db_model.save_all(new_symbole_models)

    def get_all_category_metas(self, as_map=False, type=None, page_num=1, page_size=None) -> Union[
        List[CategoryMeta], Dict[str, CategoryMeta]]:
        """
        返回所有的行业分类。
        :param as_map: 是否返回map对象
        :param type:行业类型：0未定义，  1,表示大盘指数（上证指数，深证），2 表示行业大分类  3、表示行业小分类。
        :param page_num: 从1开始
        :param page_size: 分页尺寸，为None时表示不分页。
        :return:
        """
        if page_num < 1:
            raise RuntimeError("page_num must > 0")
        by_page = page_num > 1 or page_size is not None
        with self.source.use_peewee() as db_model:
            if by_page:
                if type is not None:
                    s = (
                        db_model.select()
                            .where(
                            (db_model.meta_type == META_TYPE_CATEGORY)
                            &(db_model.type == type)
                        ).order_by(db_model.code).paginate(page_num, page_size)
                    )
                else:
                    s = (
                        db_model.select()
                            .where(
                            (db_model.meta_type == META_TYPE_CATEGORY)
                        ).order_by(db_model.code).paginate(page_num, page_size)
                    )
            else:
                if type is not None:
                    s = (
                        db_model.select()
                            .where(
                            (db_model.meta_type == META_TYPE_CATEGORY)
                            & (db_model.type == type)
                        )
                    )
                else:
                    s = (
                        db_model.select()
                            .where(
                            (db_model.meta_type == META_TYPE_CATEGORY)
                        )
                    )
            ret = {} if as_map else []
            for model in s:
                meta = model.to_meta()
                if as_map:
                    ret[meta.code] = meta
                else:
                    ret.append(meta)
            return ret

    def get_all_symbol_metas(self, as_map=False, page_num=1, page_size=None) -> Union[
        List[SymbolMeta], Dict[str, SymbolMeta]]:
        """
        返回所有的个股结构信息。
        :param as_map: 是否返回map对象
        :param page_num: 从1开始
        :param page_size: 分页尺寸，为None时表示不分页。
        :return:
        """
        if page_num < 1:
            raise RuntimeError("page_num must > 0")
        by_page = page_num > 1 or page_size is not None
        with self.source.use_peewee() as db_model:
            if by_page:
                s = (
                    db_model.select()
                        .where(
                        (db_model.meta_type == META_TYPE_SYMBOL)
                    ).order_by(db_model.code).paginate(page_num, page_size)
                )
            else:
                s = (
                    db_model.select()
                        .where(
                        (db_model.meta_type == META_TYPE_SYMBOL)
                    )
                )
            ret = {} if as_map else []
            for model in s:
                meta = model.to_meta()
                if as_map:
                    ret[meta.code] = meta
                else:
                    ret.append(meta)
            return ret

    def get_symbol_metas(self,codes:List[str],as_map=False,)-> Union[List[SymbolMeta], Dict[str, SymbolMeta]]:
        if not codes:
            return {} if as_map else []
        with self.source.use_peewee() as db_model:
            s = (
                db_model.select()
                    .where(
                    (db_model.meta_type == META_TYPE_SYMBOL)
                    & (db_model.code.in_(codes))
                )
            )
            ret = {} if as_map else []
            for model in s:
                meta = model.to_meta()
                if as_map:
                    ret[meta.code] = meta
                else:
                    ret.append(meta)
        return ret

    def get_symbol_meta(self, code: str) -> SymbolMeta:
        with self.source.use_peewee() as db_model:
            s = (
                db_model.select()
                    .where(
                    (db_model.meta_type == META_TYPE_SYMBOL)
                    & (db_model.code == code)
                ).first()
            )
            if s:
                return s.to_meta()
        return None

    def get_all_symbols(self) -> List[str]:
        """
        返回所有的个股结构代码。
        """
        with self.source.use_peewee() as db_model:
            s = (
                db_model.select(db_model.code)
                    .where(
                    (db_model.meta_type == META_TYPE_SYMBOL)
                ).order_by(db_model.code)
            )
            return [ model.code  for model in s]

    def get_category_metas_by_symbol(self, symbol, type=None) -> List[CategoryMeta]:
        """
        返回个股所属分类行业信息
        :param symbol:
        :param type:行业类型：0未定义，  1,表示大盘指数（上证指数，深证），2 表示行业大分类  3、表示行业小分类。
        :return:
        """
        with self.source.use_peewee() as db_model:
            category_code_cursor = (
                db_model.select(db_model.category_code)
                    .where(
                    (db_model.meta_type == META_TYPE_CATEGORY_SYMBOL)
                    & (db_model.code == symbol)
                )
            )
            category_codes = [ model.category_code  for model in category_code_cursor]
            if not category_codes:
                return []
            if type is not None:
                category_cursor = (
                    db_model.select()
                        .where(
                        (db_model.meta_type == META_TYPE_CATEGORY)
                        & (db_model.code.in_(category_codes))
                        &(db_model.type == type)
                    )
                )
            else:
                category_cursor = (
                    db_model.select()
                        .where(
                        (db_model.meta_type == META_TYPE_CATEGORY)
                        & (db_model.code.in_(category_codes))
                    )
                )
            return  [ model.to_meta()  for model in category_cursor]

    def get_category_meta(self, code:Union[str,List[str]],as_map = False) -> Union[CategoryMeta,List[CategoryMeta]]:
        """
        返回分类行业信息
        :param code: 行业分类代码
        :return:
        """
        if isinstance(code,str):
            with self.source.use_peewee() as db_model:
                s = (
                    db_model.select()
                        .where(
                        (db_model.meta_type == META_TYPE_CATEGORY)
                        & (db_model.code == code)
                    ).first()
                )
                if s:
                    return s.to_meta()
                return None
        else:
            with self.source.use_peewee() as db_model:
                s = (
                    db_model.select()
                        .where(
                        (db_model.meta_type == META_TYPE_CATEGORY)
                        & (db_model.code.in_(code))
                    )
                )
                return [ item.to_meta() for item in s]

    def get_category_symbol_metas(self,category_code:str,as_map = False)->Union[List[CategorySymbolMeta],Dict[str,CategorySymbolMeta]]:
        """
        返回分类下的所有个股信息
        :param category_code:
        :return:
        """
        with self.source.use_peewee() as db_model:
            s = (
                db_model.select()
                    .where(
                    (db_model.meta_type == META_TYPE_CATEGORY_SYMBOL)
                    &(db_model.category_code == category_code)
                ).order_by(db_model.code)
            )
            if as_map:
                ret = {}
                for model in s:
                    meta = model.to_meta()
                    ret[meta.code] = meta
                return ret
            return [model.to_meta() for model in s]


    def category_count(self, type=None):
        """
        市场行业分类总数
        :param type:行业类型：0未定义，  1,表示大盘指数（上证指数，深证），2 表示行业大分类  3、表示行业小分类。
        :return:
        """
        with self.source.use_peewee() as db_model:
            if type is not None:
                return db_model.select(db_model.code).where(
                    (db_model.meta_type == META_TYPE_CATEGORY)
                    &(db_model.type == type)
                ).count()
            else:
                return db_model.select(db_model.code).where(
                    (db_model.meta_type == META_TYPE_CATEGORY)

                ).count()

    def symbol_count(self):
        """
        市场个股总数
        :return:
        """
        with self.source.use_peewee() as db_model:
            return db_model.select(db_model.code).where(
                (db_model.meta_type == META_TYPE_SYMBOL)
            ).count()

    def clear(self):
        with self.source.use_peewee() as db_model:
            db_model.drop_table()
            db_model.create_table()


    def size(self)->int:
        with self.source.use_peewee() as db_model:
            return db_model.select(db_model.code).count()


    def _get_all_category_symbols(self,) -> List[CategorySymbolMeta]:
        """
        :return:
        """
        with self.source.use_peewee() as db_model:
            s = (
                db_model.select()
                    .where(
                    (db_model.meta_type == META_TYPE_CATEGORY_SYMBOL)
                ).order_by(db_model.code)
            )
            return [model.to_meta() for model in s]

